/*
 * Copyright 2021 The Closure Compiler Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.javascript.jscomp.serialization;

import static com.google.common.base.Preconditions.checkNotNull;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.javascript.jscomp.colors.Color;
import com.google.javascript.jscomp.colors.ColorId;
import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.Map.Entry;
import java.util.function.Function;
import java.util.function.Predicate;
import org.jspecify.annotations.Nullable;

/**
 * Serializes `Color`s and information about them into a `TypePool`.
 *
 * <p>Client code should call `addColor()` for each of the colors it wants to have included in the
 * `TypePool`, then call `generateTypePool()` to get the result.
 */
class ColorSerializer {
  private final SerializationOptions serializationMode;
  /**
   * Used to find the index at which a `String` will appear in the `StringPoolProto` that will
   * eventually be serialized with the `TypePool` generated by this `ColorSerializer`.
   */
  private final Function<String, Integer> getStringPoolIndexFn;
  /** Only property names that pass this predicate will be serialized. */
  private final Predicate<String> propertyFilter;

  /**
   * Stores the `Integer` values assigned to `Color`s as they are added for serialization, so they
   * can be looked up and used for references between the output `TypeProto`s.
   */
  private final LinkedHashMap<ColorId, Integer> colorIdToTypePointer = new LinkedHashMap<>();

  /**
   * Stores the `Color`s to be serialized in the order they will be serialized.
   *
   * <p>This is used as a worklist when the `TypePool` is generated.
   */
  private final ArrayList<Color> colorsInSerializedOrder = new ArrayList<>();

  /**
   * Create a ColorSerializer.
   *
   * @param serializationMode determines what parts of the `TypePool` proto to fill in.
   * @param stringPoolIndexFn Used to request an integer index to encode in place of string value.
   *     The `TypePool` generated by this object will use these indices, so a `StringPoolProto` that
   *     contains the actual `String` values needs to be serialized along with it.
   * @param propertyFilter Property names for which this returns `false` will not be included in the
   *     generated `TypeProto`s.
   */
  ColorSerializer(
      SerializationOptions serializationMode,
      Function<String, Integer> stringPoolIndexFn,
      Predicate<String> propertyFilter) {
    this.serializationMode = serializationMode;
    this.getStringPoolIndexFn = stringPoolIndexFn;
    this.propertyFilter = propertyFilter;
    // We must pre-populate the type pointers with the few axiomatic colors that won't actually
    // be serialized. These types are required to get the first few offsets in the order they
    // are specified by `TypePointers.OFFSET_TO_AXIOMATIC_COLOR`.
    for (Color color : TypePointers.OFFSET_TO_AXIOMATIC_COLOR) {
      addColor(color);
    }
  }

  /**
   * Add a collection of `Color`s to the list of those that must be serialized.
   *
   * @param colors to be serialized
   * @return list of `Integer`s that will refer to the input `Color`s in the TypePool` that this
   *     object will create. The order will match the order of the input `Color`s.
   */
  ImmutableList<Integer> addColors(Collection<Color> colors) {
    final ImmutableList.Builder<Integer> builder = ImmutableList.builder();
    for (Color color : colors) {
      builder.add(addColor(color));
    }
    return builder.build();
  }

  /**
   * Add `color` to the list of those that must be serialized (if it wasn't already there) and
   * return the `Integer` value that will refer to it in the `TypePool` that this object will
   * create.
   */
  int addColor(Color color) {
    return colorIdToTypePointer.computeIfAbsent(
        color.getId(),
        (unusedKey) -> {
          final int index = colorsInSerializedOrder.size();
          colorsInSerializedOrder.add(color);
          return index;
        });
  }

  /**
   * Generate a `TypePool` proto built from the previously added `Color`s and the arguments supplied
   * to this method.
   *
   * @param getDisambiguationSupertypesFn Given a `Color` return a set of `Color`s it inherits from.
   * @param getMismatchSourceRefsFn May be `null` if this `serializationMode` does not include debug
   *     info. Otherwise, this function must provide a set of all the source reference strings
   *     indicating code locations where the given `Color` encountered a type mismatch.
   * @return a new `TypePool` proto
   */
  TypePool generateTypePool(
      Function<Color, ImmutableSet<Color>> getDisambiguationSupertypesFn,
      @Nullable Function<Color, ImmutableSet<String>> getMismatchSourceRefsFn) {
    final TypePool.Builder typePoolBuilder = TypePool.newBuilder();
    // We use an indexed loop here for 2 reasons.
    // 1. We must skip serialization of the axiomatic colors that start our list of serialized
    //    colors.
    // 2. The logic in the loop may end up adding more colors to the list. It's often hard to tell
    //    what effect changing an iterable will have on an iteration that is in progress.
    for (int i = TypePointers.untrimOffset(0); i < colorsInSerializedOrder.size(); i++) {
      final Color color = colorsInSerializedOrder.get(i);
      final Integer typePointer = colorIdToTypePointer.get(color.getId());
      typePoolBuilder.addType(generateTypeProto(color));
      for (Color supertype : getDisambiguationSupertypesFn.apply(color)) {
        typePoolBuilder
            .addDisambiguationEdgesBuilder()
            .setSubtype(typePointer)
            .setSupertype(addColor(supertype));
      }
    }

    if (serializationMode.getIncludeDebugInfo()) {
      checkNotNull(getMismatchSourceRefsFn);
      final TypePool.DebugInfo.Builder debugInfoBuilder = typePoolBuilder.getDebugInfoBuilder();

      // Construct a map from source reference string to type pointers,
      // because that's the way the Mismatch protos work.
      // Construct entries only for those colors that we have actually serialized in order to save
      // space.
      final LinkedHashMap<String, ArrayList<Integer>> srcRefToTypePointerList =
          new LinkedHashMap<>();
      for (Color color : colorsInSerializedOrder) {
        final Integer typePointer = colorIdToTypePointer.get(color.getId());
        for (String srcRef : getMismatchSourceRefsFn.apply(color)) {
          final ArrayList<Integer> typePointerList =
              srcRefToTypePointerList.computeIfAbsent(srcRef, (key) -> new ArrayList<>());
          typePointerList.add(typePointer);
        }
      }

      // Now use the map to build the Mismatch protos and put them into the debug info.
      for (Entry<String, ArrayList<Integer>> entry : srcRefToTypePointerList.entrySet()) {
        debugInfoBuilder
            .addMismatchBuilder()
            .setSourceRef(entry.getKey())
            .addAllInvolvedColor(entry.getValue());
      }
    }
    return typePoolBuilder.build();
  }

  private TypeProto generateTypeProto(Color color) {
    final TypeProto.Builder typeProtoBuilder = TypeProto.newBuilder();
    if (color.isUnion()) {
      typeProtoBuilder.getUnionBuilder().addAllUnionMember(addColors(color.getUnionElements()));
    } else {
      final ObjectTypeProto.Builder objectTypeProtoBuilder = typeProtoBuilder.getObjectBuilder();
      objectTypeProtoBuilder
          .setIsInvalidating(color.isInvalidating())
          .setUuid(color.getId().asByteString())
          .setPropertiesKeepOriginalName(color.getPropertiesKeepOriginalName())
          .addAllInstanceType(addColors(color.getInstanceColors()))
          .addAllPrototype(addColors(color.getPrototypes()))
          .setMarkedConstructor(color.isConstructor())
          .addAllOwnProperty(getOwnPropertyStringPoolOffsets(color))
          .setClosureAssert(color.isClosureAssert());
    }
    return typeProtoBuilder.build();
  }

  private ImmutableList<Integer> getOwnPropertyStringPoolOffsets(Color color) {
    final ImmutableList.Builder<Integer> builder = ImmutableList.builder();
    for (String ownProperty : color.getOwnProperties()) {
      // The client code may know that some properties are unused in the AST, so there's no need
      // to serialize them.
      if (propertyFilter.test(ownProperty)) {
        builder.add(getStringPoolIndexFn.apply(ownProperty));
      }
    }
    return builder.build();
  }
}
